#pragma once

#include<iostream>
#include<cassert>
using namespace std;



template<class K, class V>
struct AVLTreeNode
{
    AVLTreeNode(const pair<K, V>& kv)
        :_kv(kv)
        ,_left(nullptr)
        ,_right(nullptr)
        ,_parent(nullptr)
        ,_bf(0)
    {}

    pair<K, V> _kv;
    AVLTreeNode<K, V>* _left;
    AVLTreeNode<K, V>* _right;
    AVLTreeNode<K, V>* _parent;
    int _bf;//平衡因子 -- 右子树-左子树的高度差
};

template<class K, class V>
class AVLTree
{
    typedef AVLTreeNode<K, V> Node;

private:
    Node* _root = nullptr;

public:
    bool Insert(const pair<K, V>& kv)
    {
        //1.按搜索树的规则插入
        //2.看是否违反平衡规则, 如果违反就旋转处理

        if(_root == nullptr)
        {
            _root = new Node(kv);
            _root->_bf = 0;
            return true;
        }
        Node* parent = nullptr;
        Node* cur = _root;
        while(cur)
        {
            if(cur->_kv.first < kv.first)//值
            {
                parent = cur;
                cur = cur->_right;
            }
            else if(cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }

        cur = new Node(kv);
        if(parent->_kv.first < kv.first)
            parent->_right = cur;
        else
            parent->_left = cur;
        cur->_parent = parent;

        // //调节平衡因子 -- 沿着祖先路径更新
        // if(cur == parent->_right)
        //     parent->_bf++;
        // else
        //     parent->_bf--;

        //更新平衡因子
        while(parent)//最远要更新到根
        {
            if(cur == parent->_right)
                parent->_bf++;
            else   
                parent->_bf--;
            
            //是否继续更新
            if(parent->_bf == 0)//1/-1->0 -- 插入节点到矮的那边
            {
                break;//高度不变, 更新结束
            }
            else if(parent->_bf == 1 || parent->_bf == -1)//0->1/-1 -- 插入节点导致一边变高了
            {
                //子树的高度变了, 继续更新祖先
                cur = cur->_parent;
                parent = parent->_parent;
            }
            else if(parent->_bf == 2 || parent->_bf == -2)//1/-1->2/-2 -- 插入节点导致本来高的一边又变高了
            {
                //子树不平衡, 需要旋转处理
                if(parent->_bf == 2 && cur->_bf == 1)//右边高--左单旋
                    RotateL(parent);

                else if(parent->_bf == -2 && cur->_bf == -1)//左边高--右单旋
                    RotateR(parent);

                else if(parent->_bf == -2 && cur->_bf == 1)//左右单旋
                    RotateLR(parent);
            }
        }

    }

private:
    //左单旋
    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;

        parent->_right = subRL;
        if(subRL)
            subRL->_parent = parent;

        Node* ppNode = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;

        if(parent == _root)
        {
            _root = subR;
            _root->_parent = nullptr;
        }
        else
        {
            if(parent == ppNode->_left)
                ppNode->_left = subR;
            else
                ppNode->_right = subR;
            subR->_parent = ppNode;
        }

        //更新平衡因子
        parent->_bf = 0;
        subR->_bf = 0;
    }

    //右单旋
    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;

        parent->_left = subLR;
        if(subLR)
            subLR->_parent = parent;
        
        Node* ppNode = parent->_parent;//可能是局部的子树

        subL->_right = parent;
        parent->_parent = subL;
        if(parent == _root)
        {
            _root = subL;
            _root->_parent = nullptr;
        }
        else
        {
            if(ppNode->_left == parent)
            {
                ppNode->_left = subL;
            }
            else
            {
                ppNode->_right = subL;
            }
            subL->_parent = ppNode;
        }

        //更新平衡因子
        parent->_bf = 0;
        subL->_bf = 0;
    }

    //左右单旋
    void RotateLR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;

        int bf = subLR->_bf;

        RotateL(parent->_left);
        RotateR(parent);

        //更新平衡因子
        if(bf == 0)
        {
            parent->_bf = 0;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else if(bf == 1)
        {
            parent->_bf = 0;
            subL->_bf = -1;
            subLR->_bf = 0;
        }
        else if(bf == -1)
        {
            parent->_bf = 1;
            subL->_bf = 0;
            subLR->_bf = 0;
        }
        else
        {
            //subLR->_bf 旋转前就有问题
            assert(false);
        }

    }
};